import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset
from models.stgnn import STGNN
from utils.graph import build_adjacency_matrix
from utils.feature import load_sequence_data

# Load data
X, Y = load_sequence_data('data/processed/train.csv')  # [num_samples, T, F]
X = torch.FloatTensor(X).permute(0, 2, 1).unsqueeze(-1)  # [B, F, T, 1]
Y = torch.LongTensor(Y)

# Fake user feature to build A
user_features = np.random.rand(X.shape[0], 10)
A = build_adjacency_matrix(user_features)

# Initialize model
model = STGNN(A, in_channels=X.shape[1], out_classes=3)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = torch.nn.CrossEntropyLoss()

loader = DataLoader(TensorDataset(X, Y), batch_size=64, shuffle=True)

# Train
for epoch in range(10):
    model.train()
    for xb, yb in loader:
        pred = model(xb)
        loss = loss_fn(pred.squeeze(1), yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}, loss: {loss.item():.4f}")